# This is a calcul2d, intended for creating 2-dimensional potential 
# energy surfaces of molecular systems.
# Copyright (C) 2011  Frank Uhlig
# 
# This program is free software; you can redistribute it and/or
# modify it under the terms of the GNU General Public License
# as published by the Free Software Foundation; either version 2
# of the License, or (at your option) any later version.
# 
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
# 
# You should have received a copy of the GNU General Public License
# along with this program; if not, write to the Free Software
# Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301, USA.

#!/usr/bin/env python

import os
import numpy as np
from ase.io import read
from ase.calculators.turbomole import Turbomole
from ase.optimize import FIRE
from ase.constraints import FixBondLengths
#from constants import *

class PEScalc():
    """
    """

    def __init__(self,filename=None,constraints=None,iter=None):
        """constraints is list ([i,j],float) <-- specifying bonds to constrain and float value
        """

        self.filename = filename
        self.iter = iter

        if constraints == None:
            print 'No constraints for relaxed scan specified.\n This is going to get difficult'
            quit()
        else:
            self.constraints = constraints

        self.atoms = read(self.filename)

        for constraint,length in self.constraints:
            PEScalc.set_bond_length(self,constraint,length)

        self.numerals = []
        for constraint, lengt in self.constraints:
            self.numerals.append(length)

#        i = 0
#        if self.iter is not None:
#            for constraint, length in self.constraints:
#                for stop, step in self.iter[0]:
#                    for blength in range(int(10*length),int(10*stop),int(10*step)):
#                        os.mkdir('%s-w' %str(blength))
#                        os.system('cp coord basis auxbasis mos alpha beta control %s-w' %str(blength))
#                        os.chdir('%s-w' %str(blength))
#                        PEScalc.set_bond_length(self,constraint,blength/10.)
#                        self.atoms = PEScalc.optimize(self)
#                        os.chdir('../')
#                i += 1
#        else:
#            PEScalc.optimize(self)

        if self.iter is not None:
            if len(self.iter) == 2:
                for alength in range(int(10*self.constraints[0][1]),int(10*self.iter[0][0]),int(10*self.iter[0][1])):
                    PEScalc.set_bond_length(self,constraints[0][0],alength/10.)
                    for blength in range(int(10*self.constraints[1][1]),int(10*self.iter[1][0]),int(10*self.iter[1][1])):
                        print alength, blength
                        os.mkdir('%s-w-%s' %(str(alength),str(blength)))
                        os.system('cp coord basis auxbasis mos alpha beta control %s-w-%s' %(str(alength),str(blength)))
                        os.chdir('%s-w-%s' %(str(alength),str(blength)))
                        PEScalc.set_bond_length(self,constraints[1][0],blength/10.)
                        #PEScalc.set_bond_length(self,constraints,blength/10.)
                        self.atoms = PEScalc.optimize(self)
                        os.chdir('../')
            elif len(self.iter) == 1:
                for blength in range(int(10*self.constraints[0][1]),int(10*self.iter[0][0]),int(10*self.iter[0][1])):
                    print blength
                    os.mkdir('%s-w' %(str(blength)))
                    os.system('cp coord basis auxbasis mos alpha beta control %s-w' %(str(blength)))
                    os.chdir('%s-w' %(str(blength)))
                    PEScalc.set_bond_length(self,constraints[0][0],blength/10.)
                    #PEScalc.set_bond_length(self,constraints,blength/10.)
                    self.atoms = PEScalc.optimize(self)
                    os.chdir('../')
        else:
            PEScalc.optimize(self)


    def set_bond_length(self,bond,length):
        """
        adjust molecular geometry to initially fulfill bonds
        """

        self.atoms.positions = self.atoms.positions - self.atoms.positions[bond[0],:]
        self.atoms.positions[bond[1],:] = self.atoms.positions[bond[1],:]/np.sqrt(np.sum(self.atoms.positions[bond[1],:]**2))*length

        print self.atoms.positions[bond[1],:]
        
    def optimize(self):
        """ documentation to be here soon"""

        bonds = []
        for constraint,length in self.constraints:
            bonds.append(constraint)

        print bonds

        initial = self.atoms
        initial.set_calculator(Turbomole())

        #c1 = (3, 4)
        #c2 = (7, 10)
        #c = FixBondLengths((c1, c2))
        c = FixBondLengths((bonds))
        initial.set_constraint(c)

        dyn = FIRE(initial, trajectory='geoopt.traj')
        dyn.run(fmax=0.004)

        return read('geoopt.traj', index=-1)


